// Copyright (c) 2025 Inochi Amaoto

/* SPDX-License-Identifier: GPL-2.0 */

#ifndef __VIRTIO_VAPP_H__
#define __VIRTIO_VAPP_H__

#include <linux/miscdevice.h>
#include <linux/virtio.h>
#include <linux/rbtree.h>
#include <linux/mutex.h>
#include <linux/wait.h>
#include <linux/spinlock.h>
#include <linux/hashtable.h>
#include "vapp-common.h"

#ifndef VIRTIO_ID_VAPP
#define VIRTIO_ID_VAPP 60
#endif

#define VIRTIO_VAPP_DEVICE_MINOR MISC_DYNAMIC_MINOR
#define VIRTIO_VAPP_DEVICE_MODE (umode_t)(S_IRUGO | S_IWUGO)

#define VIRTIO_VAPP_DTQ_REFILL_DIV_FACTOR   2

struct vio_vapp_notify_entity {
  struct rb_node node;
  u64 id;
  u32 type;
};

struct vio_vapp_notify_queue {
  struct rb_root queue;
  wait_queue_head_t wait;
  spinlock_t lock;
};

#define vio_vapp_notify_entity_id_cmp(_x, _y) ((_x) < (_y))

static inline void vio_vapp_notify_entity_free(struct vio_vapp_notify_entity* notify)
{
  kvfree(notify);
}

static inline struct vio_vapp_notify_entity* vio_vapp_notify_entity_alloc(u64 id, u32 type)
{
  struct vio_vapp_notify_entity* notify = kmalloc(sizeof(*notify), GFP_KERNEL);
  if (!notify)
    return NULL;

  notify->id = id;
  notify->type = type;

  return notify;
}

static inline void __vio_vapp_notify_queue_remove(struct vio_vapp_notify_queue* queue, struct vio_vapp_notify_entity* notify)
{
  rb_erase(&notify->node, &queue->queue);
}

static inline void vio_vapp_notify_queue_remove(struct vio_vapp_notify_queue* queue, struct vio_vapp_notify_entity* notify)
{
  unsigned long flags;
  spin_lock_irqsave(&queue->lock, flags);
  __vio_vapp_notify_queue_remove(queue, notify);
  spin_unlock_irqrestore(&queue->lock, flags);
}

static inline struct vio_vapp_notify_entity* vio_vapp_notify_queue_get(
struct vio_vapp_notify_queue* queue, u64 id)
{
  unsigned long flags;
  struct rb_node* node = NULL;
  struct vio_vapp_notify_entity* notify;

  spin_lock_irqsave(&queue->lock, flags);

  node = queue->queue.rb_node;
  while (node) {
    notify = container_of(node, struct vio_vapp_notify_entity, node);

    if (vio_vapp_notify_entity_id_cmp(id, notify->id))
      node = node->rb_left;
    else if (vio_vapp_notify_entity_id_cmp(notify->id, id))
      node = node->rb_right;
    else {
      spin_unlock_irqrestore(&queue->lock, flags);
      return notify;
    }
  }

  spin_unlock_irqrestore(&queue->lock, flags);

  return NULL;
}

static inline void __vio_vapp_notify_queue_mount(
struct vio_vapp_notify_queue* queue, struct vio_vapp_notify_entity* notify)
{
  unsigned long flags;
  struct rb_node **new = NULL, *parent = NULL;
  struct vio_vapp_notify_entity* this;

  spin_lock_irqsave(&queue->lock, flags);

  new = &queue->queue.rb_node;
  while (*new) {
    this = container_of(*new, struct vio_vapp_notify_entity, node);
    parent = *new;

    if (vio_vapp_notify_entity_id_cmp(notify->id, this->id))
      new = &((*new)->rb_left);
    else if (vio_vapp_notify_entity_id_cmp(this->id, notify->id))
      new = &((*new)->rb_right);
    else {
      spin_unlock_irqrestore(&queue->lock, flags);
      vio_vapp_notify_entity_free(notify);
      return;
    }
  }

  rb_link_node(&notify->node, parent, new);
  rb_insert_color(&notify->node, &queue->queue);
  spin_unlock_irqrestore(&queue->lock, flags);

  return;
}

static inline void vio_vapp_notify_queue_mount(
struct vio_vapp_notify_queue* queue, struct vio_vapp_notify_entity* notify)
{
  __vio_vapp_notify_queue_mount(queue, notify);
  wake_up(&queue->wait);
}

static inline void vio_vapp_notify_queue_recovery_mount(
struct vio_vapp_notify_queue* queue, struct vio_vapp_notify_entity* notify)
{
  __vio_vapp_notify_queue_mount(queue, notify);
}

static inline struct vio_vapp_notify_entity* vio_vapp_notify_queue_wait(
struct vio_vapp_notify_queue* queue, u64 id)
{
  struct vio_vapp_notify_entity* notify;

  int ret = wait_event_interruptible(queue->wait, (notify = vio_vapp_notify_queue_get(queue, id)));
  if (ret) {
    pr_warn("wait interruptted\n");
    return ERR_PTR(ret);
  }

  vio_vapp_notify_queue_remove(queue, notify);
  return notify;
}

static inline void vio_vapp_notify_queue_init(
struct vio_vapp_notify_queue* queue)
{
  queue->queue.rb_node = NULL;
  init_waitqueue_head(&queue->wait);
  spin_lock_init(&queue->lock);
}

static inline void vio_vapp_notify_queue_destroy(
struct vio_vapp_notify_queue* queue)
{
  struct rb_node* pos;
  unsigned long flags;

  spin_lock_irqsave(&queue->lock, flags);
  for(pos = rb_first(&queue->queue); pos; pos = rb_first(&queue->queue)) {
    struct vio_vapp_notify_entity* entity = container_of(pos, struct vio_vapp_notify_entity, node);

    __vio_vapp_notify_queue_remove(queue, entity);
    vio_vapp_notify_entity_free(entity);
  }
  spin_unlock_irqrestore(&queue->lock, flags);
}

struct vio_vapp_context {
  struct hlist_node node;
  u64 id;

  struct vio_vapp_notify_queue queue;
  struct vio_mmap_entity_tree mtree;

  struct vio_vapp_entity_tree req;
  struct vio_vapp_entity_tree rsp;
};

struct virtio_vapp_dev;
struct virtio_vapp_dtq {
  struct virtqueue* vqs[2];
  struct mutex mutex[2];
  struct work_struct rx_work;
  struct work_struct tx_work;
  struct work_struct send_work;
  bool vq_run[2];

  struct vapp_send_queue* queue;

  struct virtio_vapp_dev* vadev;

  atomic_t queued_replies;
  unsigned int rx_bufs;
  unsigned int max_rx_bufs;
};

struct virtio_vapp_dev {
  struct virtio_device* vdev;
  struct virtio_vapp_cpu_config config;

  struct vapp_send_queue cqueue;
  struct virtio_vapp_dtq cdtq;

  struct vapp_send_queue dqueue;
  struct virtio_vapp_dtq* ddtqs;

  atomic_t id_pool;
};

int vio_vapp_init(void);
void vio_vapp_exit(void);
struct vio_vapp_context* virtio_vapp_get_vio_dev(u64 fp);
struct virtio_vapp_dev* virtio_vapp_get_dev_rcu(void);

static inline u64 virtio_vapp_get_dev_id(void)
{
  u64 id = 0;
  struct virtio_vapp_dev* dev;
  rcu_read_lock();

  dev = virtio_vapp_get_dev_rcu();
  if (!dev)
    goto empty;

 again:
  id = atomic_fetch_add(1, &dev->id_pool);
  if (id == 0)
    goto again;

 empty:
  rcu_read_unlock();

  return id;
}

static inline struct vio_vapp_entity_tree* virtio_vio_vapp_select_tree(struct vio_vapp_context* vapp, u32 type)
{
  if (type & VIO_TYPE_GUEST)
    return &vapp->req;
  else
    return &vapp->rsp;
}

static inline struct vio_mmap_entity_tree* vio_mmap_entity_tree_from_file(struct file* fp)
{
  struct vio_vapp_context* ctx = fp->private_data;
  return &ctx->mtree;
}

#endif // ! __VIRTIO_VAPP_H__

